{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Corrective RAG Pack\n",
    "\n",
    "This notebook walks through using the `CorrectiveRAGPack` based on the paper [Corrective Retrieval Augmented Generation](https://arxiv.org/abs/2401.15884).\n",
    "\n",
    "A brief understanding of the paper:\n",
    "\n",
    "\n",
    "Corrective Retrieval Augmented Generation (CRAG) is a method designed to enhance the robustness of language model generation by evaluating and augmenting the relevance of retrieved documents through a an evaluator and large-scale web searches, ensuring more accurate and reliable information is used in generation.\n",
    "\n",
    "We use `GPT-4` as a relevancy evaluator and `Tavily AI` for web searches. So, we recommend getting `OPENAI_API_KEY` and `tavily_ai_api_key` before proceeding further."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai llama-index-tools-tavily-research llama-index-embeddings-openai\n",
    "%pip install llama-index-packs-corrective-rag\n",
    "%pip install llama-index-readers-file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"YOUR OPENAI API KEY\"\n",
    "\n",
    "tavily_ai_api_key = \"<tavily_ai_api_key>\"\n",
    "\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download `Llama2` paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100 13.0M  100 13.0M    0     0  7397k      0  0:00:01  0:00:01 --:--:-- 7415k\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p 'data/'\n",
    "!curl 'https://arxiv.org/pdf/2307.09288.pdf' -o 'data/llama2.pdf'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(\"data\").load_data()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the CorrectiveRAGPack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.packs.corrective_rag import CorrectiveRAGPack\n",
    "\n",
    "corrective_rag_pack = CorrectiveRAGPack(documents, tavily_ai_apikey=tavily_ai_api_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "Llama2 was pretrained using an optimized auto-regressive transformer. Several changes were made to improve performance, including more robust data cleaning, updated data mixes, training on 40% more total tokens, doubling the context length, and using grouped-query attention (GQA) to improve inference scalability for larger models."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Markdown, display\n",
    "\n",
    "response = corrective_rag_pack.run(\"How was Llama2 pretrained?\")\n",
    "display(Markdown(str(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The latest ChatGPT memory update allows the chatbot to carry what it learns between chats, enabling it to provide more relevant responses. Users can ask ChatGPT to remember specific information or let it pick up details itself. This memory feature allows ChatGPT to remember and forget things based on user input, enhancing user interaction with personalized continuity and emotional intelligence. Users can toggle this function on or off within the ChatGPT settings menu, and when disabled, ChatGPT will neither generate nor access memories."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Markdown, display\n",
    "\n",
    "response = corrective_rag_pack.run(\n",
    "    \"What is the functionality of latest ChatGPT memory.\"\n",
    ")\n",
    "display(Markdown(str(response)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "corrective-rag",
   "language": "python",
   "name": "corrective-rag"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  },
  "vscode": {
   "interpreter": {
    "hash": "b1d2a638b53f4d7129cb7686d8e3b97ae1d80a593a1618479f60cef5591ea888"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
